# Copyright 2024 CHINA MERCHANTS BANK CO., LTD.
# Copyright 2024 Huawei Technologies Co., Ltd
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

""" run eval """
import os
import argparse
from mindformers.tools.utils import str2bool
from easydict import EasyDict

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--models", type=str, required=True)
    parser.add_argument("--datasets", type=str, required=True)
    parser.add_argument("--summarizer", type=str, default=None)
    parser.add_argument("--debug", type=str2bool, default=False)
    parser.add_argument("--dry-run", type=str2bool, default=False)
    parser.add_argument("--mode", type=str, default=None)
    parser.add_argument("--reuse", type=str, default=None)
    parser.add_argument("--work-dir", type=str, default=None)
    parser.add_argument("--config-dir", type=str, default=None)
    parser.add_argument("--max-num-workers", type=int, default=None)
    parser.add_argument("--dump-eval-details", type=str2bool, default=False)
    parser.add_argument("--dump-extract-rate", type=str2bool, default=False)
    parser.add_argument("--merge-ds", type=str2bool, default=False)
    parser.add_argument("--disable-cb", type=str2bool, default=False)
    parser.add_argument("--num-prompts", type=int, default=None)
    args_ = parser.parse_args()
    args = EasyDict(vars(args_))
    return args

def main():
    cmd = "ais_bench"
    args = get_args()
    for key, value in args.items():
        key = key.replace('_', '-')
        if value:
            if key in ['debug', 'dry-run', 'dump-eval-details', 'dump-extract-rate', 'merge-ds', 'num-prompts']:
                cmd += f" --{key}"
            else:
                cmd += f" --{key} {value}"

    print('run cmd: ', cmd)
    os.system(cmd)

if __name__ == '__main__':
    main()
